from multiprocessing import pool
import zipfile
import os
from PIL import Image
from shutil import copyfile
from multiprocessing.dummy import Pool as Pool, Lock
import sys


class Tile_Tools:
    def __init__(self) -> None:
        self.lock = Lock()
        self.pool = Pool()
        self.count = 1

    def clear_function(self, input_dir):
        """
            清洗目录形式的PNG切片，删除无数据(rgba(0,0,0,0))的空256*256瓦片
            @param input_dir: 所需要清洗的PNG数据目录
            @return : void
        """
        # 初始化切片清洗
        file_list = []
        print("初始化png切片清理")
        for root, dirs, files in os.walk(input_dir):
            for file in files:
                file_path = os.path.join(root, file)
                file_size = os.path.getsize(file_path)
                if file_size <= 1246 and file_path.split(".")[-1] == "png":
                    file_list.append(file_path.replace("\\", "/"))
                elif file_size <= 214 and file_path.split(".")[-1] == "webp":
                    file_list.append(file_path.replace("\\", "/"))
        self.length = len(file_list)

        # 多线程执行清洗切片
        self.pool.map(self.__clear_with_counter, file_list)

        # 清洗空目录
        for root, dirs, files in os.walk(input_dir):
            if not os.listdir(root):
                os.rmdir(root)

    def __clear_with_counter(self, png):
        self.lock.acquire()
        if os.path.exists(png):
            os.remove(png)
            ratio = ('%.2f' % (self.count / self.length * 100))
            sys.stdout.write("\r" + "正在删除无数据切片, 进度：" + str(ratio) + "%")
            sys.stdout.flush()
        self.count += 1
        self.lock.release()

    def single_png_to_single_webp_function(self, input_dir, output_dir):
        """
            单个png转换成webp新目录
            @param input_dir: png目录
            @param output_dir: webp新目录
            @return : void
        """
        path_list = []
        for root, dirs, files in os.walk(input_dir):
            for file in files:
                if file.split(".")[-1] == "png":
                    png_file_path = os.path.join(root, file).replace("\\", "/")
                    webp_file_path = png_file_path.replace(".png", ".webp").replace(
                        input_dir.replace("\\", "/"), output_dir.replace("\\", "/"))
                    path_dict = {
                        "png": png_file_path,
                        "webp": webp_file_path
                    }
                    path_list.append(path_dict)

        self.length = len(path_list)
        self.pool.map(self.__png_convert_webp_with_counter, path_list)

    def __png_convert_webp_with_counter(self, path_dict):
        png_path = path_dict["png"]
        webp_path = path_dict["webp"]
        webp_file_name = webp_path.split("/")[-1]
        webp_root = webp_path.replace("/" + webp_file_name, "")
        if not os.path.exists(webp_root):
            os.makedirs(webp_root)
        img = Image.open(png_path)
        img.save(webp_path, "WEBP")
        img.close()
        self.lock.acquire()
        self.count += 1
        ratio = ('%.2f' % (self.count / self.length * 100))
        self.lock.release()
        sys.stdout.write("\r" + "正在生成webp切片, 进度：" + str(ratio) + "%")
        sys.stdout.flush()

    def single_png_mergen_function(self, input_dir, output_dir):
        """
            多个png总目录合并成一个新目录，同名瓦片做融合处理
            @param input_dir: png总目录
            @param output_dir: png合并后新目录
            @return : void
        """
        file_list = []
        for root, dirs, files in os.walk(input_dir):
            for file in files:
                if file.split(".")[-1] == "png":
                    input_file_path = os.path.join(root, file)
                    tiles_y = input_file_path.split("\\")[-1]
                    tiles_x = input_file_path.split("\\")[-2]
                    tiles_z = input_file_path.split("\\")[-3]
                    output_dir_path = os.path.join(
                        output_dir, tiles_z, tiles_x)
                    output_file_path = os.path.join(
                        output_dir, tiles_z, tiles_x, tiles_y)

                    path_dict = {
                        "input_file_path": input_file_path.replace("\\", "/"),
                        "output_file_path": output_file_path.replace("\\", "/"),
                        "output_dir_path": output_dir_path.replace("\\", "/")
                    }
                    file_list.append(path_dict)

        self.length = len(file_list)

        self.pool.map(self.__png_mergen_with_counter, file_list)

    def __png_mergen_with_counter(self, path_dict):
        if not os.path.exists(path_dict["output_dir_path"]):
            os.makedirs(path_dict["output_dir_path"])

        if not os.path.exists(path_dict["output_file_path"]):
            try:
                copyfile(path_dict["input_file_path"],
                         path_dict["output_file_path"])
                self.lock.acquire()
                self.count += 1
                ratio = ('%.2f' % (self.count / self.length * 100))
                self.lock.release()
                sys.stdout.write("\r" + "正在合并png切片, 进度：" + str(ratio) + "%")
                sys.stdout.flush()
            except IOError as e:
                pass
        else:
            im = Image.open(path_dict["output_file_path"]).convert("RGBA")
            im_tar = Image.open(path_dict["input_file_path"]).convert("RGBA")
            im.paste(im_tar, mask=im_tar)
            im.save(path_dict["output_file_path"])
            im.close()
            im_tar.close()
            self.lock.acquire()
            self.count += 1
            ratio = ('%.2f' % (self.count / self.length * 100))
            self.lock.release()
            sys.stdout.write("\r" + "正在合并png切片, 进度：" + str(ratio) + "%")
            sys.stdout.flush()

    def single_webp_mergen_function(self, input_dir, output_dir):
        """
            多个webp总目录合并成一个新目录，同名瓦片做融合处理
            @param input_dir: webp总目录
            @param output_dir: webp合并后新目录
            @return : void
        """
        file_list = []
        for root, dirs, files in os.walk(input_dir):
            for file in files:
                if file.split(".")[-1] == "webp":
                    input_file_path = os.path.join(root, file)
                    tiles_y = input_file_path.split("\\")[-1]
                    tiles_x = input_file_path.split("\\")[-2]
                    tiles_z = input_file_path.split("\\")[-3]
                    output_dir_path = os.path.join(
                        output_dir, tiles_z, tiles_x)
                    output_file_path = os.path.join(
                        output_dir, tiles_z, tiles_x, tiles_y)

                    path_dict = {
                        "input_file_path": input_file_path.replace("\\", "/"),
                        "output_file_path": output_file_path.replace("\\", "/"),
                        "output_dir_path": output_dir_path.replace("\\", "/")
                    }
                    file_list.append(path_dict)

        self.length = len(file_list)

        self.pool.map(self.__webp_mergen_with_counter, file_list)

    def __webp_mergen_with_counter(self, path_dict):
        if not os.path.exists(path_dict["output_dir_path"]):
            os.makedirs(path_dict["output_dir_path"])

        if not os.path.exists(path_dict["output_file_path"]):
            try:
                copyfile(path_dict["input_file_path"],
                         path_dict["output_file_path"])
                self.lock.acquire()
                self.count += 1
                ratio = ('%.2f' % (self.count / self.length * 100))
                self.lock.release()
                sys.stdout.write("\r" + "正在合并webp切片, 进度：" + str(ratio) + "%")
                sys.stdout.flush()
            except IOError as e:
                pass
        else:
            im = Image.open(path_dict["output_file_path"]).convert("RGBA")
            im_tar = Image.open(path_dict["input_file_path"]).convert("RGBA")
            im.paste(im, mask=im_tar)
            im.save(path_dict["output_file_path"])
            im.close()
            im_tar.close()
            self.lock.acquire()
            self.count += 1
            ratio = ('%.2f' % (self.count / self.length * 100))
            self.lock.release()
            sys.stdout.write("\r" + "正在合并webp切片, 进度：" + str(ratio) + "%")
            sys.stdout.flush()

    def zip_to_webp_function(self, input_dir, output_dir):
        """
            多个png切片压缩包总目录合并成一个新目录，同名瓦片做融合处理，并将格式转换为webp
            @param input_dir: png切片压缩包总目录
            @param output_dir: png切片压缩包合并，转换为webp后新目录
            @return : void
        """
        self.input_dir = input_dir
        self.output_dir = output_dir

        for root, dirs, files in os.walk(input_dir):
            for file in files:
                zip_file_path = os.path.join(root, file)
                if zip_file_path.split(".")[-1] == "zip":
                    self.extract_file_list = []
                    self.count = 1
                    self.zip_name = "   [" + file.replace(".zip","") + "]   "
                    with zipfile.ZipFile(zip_file_path, mode="r") as zf:
                        self.zf = zf
                        sys.stdout.write("\n")
                        sys.stdout.flush()
                        self.pool.map(self.__get_name_list, zf.namelist())
                        self.length = len(self.extract_file_list)
                        self.pool.map(self.__extract_files, self.extract_file_list)
                    zf.close()
        
        file_list = []
        for root,dirs, files in os.walk(output_dir):
            for file in files:
                png_file_path = os.path.join(root, file).replace("\\", "/")
                webp_file_path = os.path.join(root, file).replace("\\", "/").replace(".png",".webp")
                file_list.append({"png": png_file_path, "webp": webp_file_path})
        
        self.count = 1
        self.length = len(file_list)
        sys.stdout.write("\n")
        sys.stdout.flush()
        self.pool.map(self.__delet_png_to_webp_with_counter, file_list)
                
        temp_file = os.path.join(input_dir, "temp")
        if os.path.exists(temp_file):
            for root, dirs,files in os.walk(temp_file):
                for file in files:
                    os.remove(os.path.join(root,file))
        
            for root, dirs, files in os.walk(temp_file):
                if not os.listdir(root):
                    os.rmdir(root)
        

            for root, dirs, files in os.walk(temp_file):
                if not os.listdir(root):
                    os.rmdir(root)
            
            os.rmdir(temp_file)

    def __get_name_list(self, name):
        if name.split(".")[-1] == "png":
            output_file_path = os.path.join(self.output_dir, name).replace("\\","/")
            extract_file_dict = {
                "name": name,
                "output_file_path": output_file_path 
            }
            self.extract_file_list.append(extract_file_dict)

    def __extract_files(self, dict):
        name = dict["name"]
        output_file_path = dict["output_file_path"]
        
        if not os.path.exists(output_file_path) and name.split(".")[-1] == "png":
            self.lock.acquire()
            self.zf.extract(name, self.output_dir)
            ratio = ('%.2f' % (self.count / self.length * 100))
            self.count += 1
            self.lock.release()
            sys.stdout.write("正在解压"+ self.zip_name +", 进度：" + str(ratio) + "%" + "\r")
            sys.stdout.flush()
        elif os.path.exists(output_file_path) and name.split(".")[-1] == "png":
            self.lock.acquire()
            self.zf.extract(name, self.input_dir + "\\temp")
            temp_file_path = os.path.join(self.input_dir, "temp", name)
            im = Image.open(output_file_path)
            im_tar = Image.open(temp_file_path)
            im.paste(im_tar, mask=im_tar)
            im.save(output_file_path)
            im.close()
            im_tar.close()
            os.remove(os.path.join(self.input_dir, "temp", name))
            ratio = ('%.2f' % (self.count / self.length * 100))
            self.count += 1
            self.lock.release()
            sys.stdout.write("正在解压"+ self.zip_name +", 进度：" + str(ratio) + "%"  + "\r")
            sys.stdout.flush()
    
        

    def __delet_png_to_webp_with_counter(self, file_dict):
        """
            单个png切片目录内的格式转换，瓦片格式转换为webp并删除同位置的png
            @param input_dir: png切片总目录
            @return : void
        """
        png_file_path = file_dict["png"]
        webp_file_path = file_dict["webp"]

        if os.path.exists(png_file_path):
            self.lock.acquire()
            img = Image.open(png_file_path)
            img.save(webp_file_path, "WEBP")
            img.close()
            os.remove(png_file_path)
            ratio = ('%.2f' % (self.count / self.length * 100))
            self.count += 1
            self.lock.release()
            sys.stdout.write("png切片格式转换中"+", 进度：" + str(ratio) + "%"  + "\r")
            sys.stdout.flush()


